import threading
from MyTT import EMA
import numpy as np
import pandas as pd
from tqsdk import TqApi
from db_manager import connect_db, get_monitored_products  # 新增导入
import time
import asyncio

class OpenCloseStrategy03:
    def __init__(self, api: TqApi, symbols, lots=1):  # 新增symbols参数
        self.api = api
        self.symbols = symbols  # 新增字段
        self.lots = lots
        self.running = False
        self.threads = []
        # 移除数据库连接（改由外部传入symbols）

    def start(self):
        """启动策略"""
        self.running = True
        # Use the same event loop as TqApi
        self.loop = self.api._loop  
        # Create tasks directly in the existing loop
        tasks = [self.loop.create_task(self._process_product(symbol)) for symbol in self.symbols]
        self.tasks = tasks

    def stop(self):
        """停止策略"""
        self.running = False
        if hasattr(self, 'tasks'):
            for task in self.tasks:
                task.cancel()

    async def _process_product(self, symbol):
        async with self.api.register_update_notify() as update_chan:
            async for _ in update_chan:
                if not self.running:
                    break
                # Remove the inner while loop completely
                # 处理数据更新逻辑
                klines = self.api.get_kline_serial(symbol, 300, 500)
                if self.api.is_changing(klines.iloc[-1], "datetime"):  # Restore original condition
                    print("系统正在监控： "+symbol+" ")
                    if not klines.close.iloc[-1]: # 如果最后一根K线的收盘价为空，跳过
                        # 计算技术指标
                        q = (3*klines.close + klines.low + klines.open + klines.high) / 6
                        ma_values = self._calculate_ma(q)
                        short_ma = ma_values['ma']
                        ema_7 = ma_values['ema']
                        
                        # 获取当前持仓
                        position = self.api.get_position(symbol)
                        
                        # 修改后的交易逻辑
                        # 平全部空仓
                        if short_ma[-2] > ema_7[-2]:  
                            if position.pos_short > 0:
                                self.api.insert_order(symbol, direction="BUY", offset="CLOSE", 
                                                    volume=position.pos_short)
                        # 开多仓
                        if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] < ema_7[-3] and short_ma[-2] > ema_7[-2]:
                            self.api.insert_order(symbol, direction="BUY", offset="OPEN", 
                                                    volume=self.lots)
                            
                        # 平全部多仓
                        if short_ma[-2] < ema_7[-2]:  
                            if position.pos_long > 0:
                                self.api.insert_order(symbol, direction="SELL", offset="CLOSE", 
                                                    volume=position.pos_long)
                        # 开空仓
                        if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] > ema_7[-3] and short_ma[-2] < ema_7[-2]:
                            self.api.insert_order(symbol, direction="SELL", offset="OPEN", 
                                                    volume=self.lots)
                            
                        # 打印交易信息
                        # print(f"当前持仓: {position.pos_long} 多头, {position.pos_short} 空头")
                        # print(f"当前K线: {klines.iloc[-1].datetime}, 收盘价: {klines.iloc[-1].close}")
    
    def _calculate_ma(self, q):
        """自定义权重计算+MyTT优化"""
      
        # 保持原有的权重计算方式
        weights = np.arange(26, 0, -1)
        ma = []
        q_array = np.array(q)
        for i in range(len(q_array)):
            if i < 26:
                ma.append(np.nan)
                continue
            window = q_array[i-26:i+1]
            ma.append(np.dot(window, weights) / 351)
        
        # 使用MyTT的EMA计算
        ema_7 = EMA(pd.Series(ma), 7)
        
        return {'ma': ma, 'ema': ema_7.tolist()}